package com.lifq.hive.udf;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFUtils;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorConverter;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.typeinfo.BaseCharTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoFactory;

/**
 * @Author: LiFuqiang
 * @Date: 2023/6/16 9:35
 * @Version 1.0
 * @description
 */
public class ToUpperCaseUDF extends GenericUDF {

    private transient PrimitiveObjectInspector argumentOI;
    private transient PrimitiveObjectInspectorConverter.StringConverter stringConverter;
    private transient PrimitiveObjectInspector.PrimitiveCategory returnType = PrimitiveObjectInspector.PrimitiveCategory.STRING;
    private transient GenericUDFUtils.StringHelper returnHelper;

    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        if (arguments.length != 1) {
            throw new UDFArgumentLengthException(
                    "UPPER requires 1 argument, got " + arguments.length);
        }

        if (arguments[0].getCategory() != ObjectInspector.Category.PRIMITIVE) {
            throw new UDFArgumentException(
                    "UPPER only takes primitive types, got " + argumentOI.getTypeName());
        }
        argumentOI = (PrimitiveObjectInspector) arguments[0];

        stringConverter = new PrimitiveObjectInspectorConverter.StringConverter(argumentOI);
        PrimitiveObjectInspector.PrimitiveCategory inputType = argumentOI.getPrimitiveCategory();
        ObjectInspector outputOI = null;
        BaseCharTypeInfo typeInfo;
        switch (inputType) {
            case CHAR:
                // return type should have same length as the input.
                returnType = inputType;
                typeInfo = TypeInfoFactory.getCharTypeInfo(
                        GenericUDFUtils.StringHelper.getFixedStringSizeForType(argumentOI));
                outputOI = PrimitiveObjectInspectorFactory.getPrimitiveWritableObjectInspector(
                        typeInfo);
                break;
            case VARCHAR:
                // return type should have same length as the input.
                returnType = inputType;
                typeInfo = TypeInfoFactory.getVarcharTypeInfo(
                        GenericUDFUtils.StringHelper.getFixedStringSizeForType(argumentOI));
                outputOI = PrimitiveObjectInspectorFactory.getPrimitiveWritableObjectInspector(
                        typeInfo);
                break;
            default:
                returnType = PrimitiveObjectInspector.PrimitiveCategory.STRING;
                outputOI = PrimitiveObjectInspectorFactory.writableStringObjectInspector;
                break;
        }
        returnHelper = new GenericUDFUtils.StringHelper(returnType);
        return outputOI;
    }

    @Override
    public Object evaluate(DeferredObject[] arguments) throws HiveException {
        String val = null;
        if (arguments[0] != null) {
            val = (String) stringConverter.convert(arguments[0].get());
        }
        if (val == null) {
            return null;
        }
        val = val.toUpperCase();
        return returnHelper.setReturnValue(val);
    }

    @Override
    public String getDisplayString(String[] children) {
        return getStandardDisplayString("upper", children);
    }
}
